import gradio as gr
import pickle
import numpy as np
from PIL import Image
import io

# 加载预训练的 KNN 模型
def load_model(model_path='optimal_knn_model.pkl'):
    """
    从保存的 pickle 文件中加载 KNN 模型
    """
    try:
        with open(model_path, 'rb') as file:
            model = pickle.load(file)
        print("模型加载成功!")
        return model
    except FileNotFoundError:
        print(f"错误: 找不到模型文件 {model_path}")
        return None
    except Exception as e:
        print(f"加载模型时出错: {e}")
        return None

# 图像预处理函数
def preprocess_image(image):
    """
    将输入图像转换为模型可接受的格式
    """
    try:
        # 如果图像是 RGBA 格式，转换为 RGB
        if image.mode == 'RGBA':
            image = image.convert('RGB')
        
        # 调整图像大小为 28x28 像素（MNIST 数据集标准尺寸）
        image = image.resize((28, 28))
        
        # 转换为灰度图
        image = image.convert('L')
        
        # 转换为 numpy 数组
        image_array = np.array(image)
        
        # 反转颜色（如果需要，因为 MNIST 是白底黑字）
        image_array = 255 - image_array
        
        # 归一化到 [0, 1] 范围
        image_array = image_array / 255.0
        
        # 展平为 1D 数组 (784 个特征)
        image_flattened = image_array.flatten().reshape(1, -1)
        
        return image_flattened
    except Exception as e:
        print(f"图像预处理出错: {e}")
        return None

# 预测函数
def predict_digit(image):
    """
    使用加载的模型进行预测
    """
    # 加载模型
    model = load_model()
    if model is None:
        return "模型加载失败，请检查模型文件是否存在"
    
    # 预处理图像
    processed_image = preprocess_image(image)
    if processed_image is None:
        return "图像处理失败，请重试"
    
    try:
        # 进行预测
        prediction = model.predict(processed_image)
        # 返回预测结果
        return f"预测数字: {prediction[0]}"
    except Exception as e:
        return f"预测过程中出错: {e}"

# 创建 Gradio 界面
def create_webapp():
    """
    创建并启动 Gradio Web 应用
    """
    # 界面描述
    description = """
    # 手写数字识别应用
    
    上传手写数字图像或使用画板绘制数字，模型将预测数字内容。
    
    **使用说明:**
    1. 在画板上绘制一个数字（0-9）
    2. 点击"提交"进行识别
    3. 查看右侧的预测结果
    
    **提示:** 尽量让数字居中且清晰，识别效果会更好。
    """
    
    # 创建界面
    iface = gr.Interface(
        fn=predict_digit,  # 预测函数
        inputs=gr.Image(type="pil", label="手写数字图像"),  # 图像输入
        outputs=gr.Textbox(label="预测结果"),  # 文本输出
        title="手写数字识别器",
        description=description,
        examples=[  # 示例图像（可选）
            # 你可以在这里添加一些示例图像路径
        ],
        live=False  # 非实时预测，需要点击提交
    )
    
    return iface

# 主程序
if __name__ == "__main__":
    print("正在启动手写数字识别 Web 应用...")
    
    # 创建并启动应用
    app = create_webapp()
    app.launch(
        server_name="0.0.0.0",  # 允许外部访问
        server_port=7861,       # 端口号
        share=False             # 是否创建公共链接
    )
# create_sample_model.py - 创建示例模型（如果需要）
from sklearn import datasets
from sklearn.model_selection import train_test_split
from sklearn.neighbors import KNeighborsClassifier
import pickle

# 加载 MNIST 数据集（简化版，实际应使用完整数据集）
digits = datasets.load_digits()
X, y = digits.data, digits.target

# 分割数据集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 创建并训练 KNN 模型
knn = KNeighborsClassifier(n_neighbors=3)
knn.fit(X_train, y_train)

# 保存模型
with open('optimal_knn_model.pkl', 'wb') as file:
    pickle.dump(knn, file)

print("示例模型已创建并保存为 optimal_knn_model.pkl")